import re
from dalchemy.parsers import DalchemyParser
# from dalchemy.parsers import DalchemyParser, register_parser
__all__ = ["QuestionAnswerJsonParser",
           ]
# @register_parser("qa_json")
class QuestionAnswerJsonParser(DalchemyParser):
    """Parse the output of an LLM .
        Example：
        input_str = '''
                     {
                        "question": "今天天气怎么样",
                        "answer": "今天天气很好。"
                    },
                    {
                        "question": "中国首都？",
                        "answer": "北京"
                    }
                    '''
        parser = QuestionAnswerJsonParser()
        res = parser.parse(input_str)
        print(res)
        >> [{'question': '今天天气怎么样？', 'answer': '今天天气很好。'},
            {'question': '中国首都？', 'answer': '北京。'},

    """
    q_symbol: str = "question"
    a_symbol: str = "answer"


    def parse(self, text: str):
        """Parse the output to qa pairs."""
        question_regex = re.compile(rf'[\'\"]{self.q_symbol}[\'\"]:[\s]?[\'\"](?P<que>.*?)[\'\"],')
        answer_regex = re.compile(rf'[\'\"]{self.a_symbol}[\'\"]:[\s]?[\'\"](?P<ans>.*?)[\'\"]')
        m_q = re.finditer(question_regex, text)
        m_a = re.finditer(answer_regex, text)
        if m_q and m_a:
            questions = [q['que'].strip() for q in m_q]
            answers = [a['ans'].strip() for a in m_a]
        assert len(questions) == len(answers), "que and ans 's length mast same!"
        qa_pairs = []
        for que, ans in zip(questions, answers):
            qa_pairs.append({self.q_symbol: que, self.a_symbol: ans})
        return qa_pairs

    def get_format_instructions(self):
        format_str = '''
        请以JSONL格式输出，包括问题question和回答answer字段【注意格式问题：如双引号、转义，并单行输出】。
        '''
        return format_str




